#include <iostream>
#include <vector>
#include <hip/hip_runtime.h>

#define CHECK_HIP(cmd) \
{\
    hipError_t error = cmd;\
    if (error != hipSuccess) {\
        std::cerr << "HIP error: " << hipGetErrorString(error) << " at line " << __LINE__ << std::endl;\
        exit(EXIT_FAILURE);\
    }\
}

// GPU 核函数：向量加法
__global__ void vectorAdd(float* A, float* B, float* C, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    // 获取 GPU 设备数量
    int numDevices = 0;
    CHECK_HIP(hipGetDeviceCount(&numDevices));
    std::cout << "Found " << numDevices << " GPU devices" << std::endl;

    if (numDevices < 2) {
        std::cerr << "Need at least 2 GPUs!" << std::endl;
        return EXIT_FAILURE;
    }

    const int N = 1 << 20;  // 1M elements
    const int bytes = N * sizeof(float);
    const int threadsPerBlock = 256;
    const int blocks = (N + threadsPerBlock - 1) / threadsPerBlock;

    // 在 GPU 0 和 GPU 1 上分配内存
    float *d_A0, *d_B0, *d_C0;  // GPU 0
    float *d_A1, *d_B1, *d_C1;  // GPU 1

    // 设置设备上下文
    CHECK_HIP(hipSetDevice(0));
    CHECK_HIP(hipMalloc(&d_A0, bytes));
    CHECK_HIP(hipMalloc(&d_B0, bytes));
    CHECK_HIP(hipMalloc(&d_C0, bytes));

    CHECK_HIP(hipSetDevice(1));
    CHECK_HIP(hipMalloc(&d_A1, bytes));
    CHECK_HIP(hipMalloc(&d_B1, bytes));
    CHECK_HIP(hipMalloc(&d_C1, bytes));

    // 初始化主机数据
    std::vector<float> h_A(N, 1.0f);  // 全1向量
    std::vector<float> h_B(N, 2.0f);  // 全2向量
    std::vector<float> h_C(N, 0.0f);

    // 将数据拷贝到 GPU 0
    CHECK_HIP(hipSetDevice(0));
    CHECK_HIP(hipMemcpy(d_A0, h_A.data(), bytes, hipMemcpyHostToDevice));
    CHECK_HIP(hipMemcpy(d_B0, h_B.data(), bytes, hipMemcpyHostToDevice));

    // 从 GPU 0 拷贝数据到 GPU 1 (跨设备拷贝)
    CHECK_HIP(hipMemcpyPeer(d_A1, 1, d_A0, 0, bytes));
    CHECK_HIP(hipMemcpyPeer(d_B1, 1, d_B0, 0, bytes));

    // 在两个 GPU 上分别执行计算
    // GPU 0 计算
    CHECK_HIP(hipSetDevice(0));
    vectorAdd<<<blocks, threadsPerBlock>>>(d_A0, d_B0, d_C0, N);

    // GPU 1 计算
    CHECK_HIP(hipSetDevice(1));
    vectorAdd<<<blocks, threadsPerBlock>>>(d_A1, d_B1, d_C1, N);

    // 将 GPU 1 的结果拷贝回 GPU 0
    CHECK_HIP(hipMemcpyPeer(d_C0, 0, d_C1, 1, bytes));

    // 将结果拷贝回主机
    CHECK_HIP(hipSetDevice(0));
    CHECK_HIP(hipMemcpy(h_C.data(), d_C0, bytes, hipMemcpyDeviceToHost));

    // 验证结果
    bool success = true;
    for (int i = 0; i < N; i++) {
        if (fabs(h_C[i] - 3.0f) > 1e-6) {  // 1 + 2 = 3
            success = false;
            break;
        }
    }
    std::cout << "Verification: " << (success ? "SUCCESS" : "FAILURE") << std::endl;

    // 释放资源
    CHECK_HIP(hipFree(d_A0));
    CHECK_HIP(hipFree(d_B0));
    CHECK_HIP(hipFree(d_C0));
    CHECK_HIP(hipFree(d_A1));
    CHECK_HIP(hipFree(d_B1));
    CHECK_HIP(hipFree(d_C1));

    return EXIT_SUCCESS;
}